import uuid
import logging
from facenet_pytorch import MTCNN, InceptionResnetV1
import torch
import time
from common.common import get_vector_index_name, base_64_encode
import os
import json
import numpy as np

from common.logger_config import setup_logger

logger = setup_logger(__name__)

device = 'cpu'
# 用于人脸检测的多任务级联网络（Multi-Task Cascaded Convolutional Networks）
mtcnn = MTCNN(
    image_size=160, margin=0, min_face_size=20,
    thresholds=[0.6, 0.7, 0.7], factor=0.709, post_process=True,
    device=device,
    keep_all=True
)
# 一个预训练的人脸识别模型，用于生成人脸特征向量（embedding）
# 加载预训练的 InceptionResnetV1 模型，使用 vggface2 数据集进行预训练，适合人脸特征提取。
resnet = InceptionResnetV1(pretrained='vggface2').eval().to(device)


class FaceResponse:
    def __init__(self, personId, box, prob, vector):
        # id
        self.personId = personId
        # box
        self.box = box
        # box
        self.prob = prob
        # 向量值
        self.vector = vector

    def to_dict(self):
        return {
            "personId": self.personId,
            "box": self.box,
            "prob": self.prob,
            "vector": self.vector
        }


class ImageFace:
    def __init__(self, prob, data, index, imageId):
        # 概率
        self.prob = prob
        # 向量值
        self.data = data
        # index
        self.index = index
        # imageId
        self.imageId = imageId

    def __lt__(self, other):
        return self.prob < other.prob

    # def __repr__(self):
    #     # return f"prob is {self.prob} with aligned {self.data}"
    #     return f"prob is {self.prob}"

    def to_dict(self):
        return {
            "prob": self.prob,
            "data": self.data,
            "index": self.index,
            "imageId": self.imageId
        }


class Person:
    def __init__(self, id, prob, data, imageId, imageIds):
        self.id = id

        # 最大的imageFace的概率
        self.prob = prob

        # 最大的imageFace的向量值
        self.data = data

        # imageId
        self.imageId = imageId

        # 这个人包含的imageFace的图片ID列表
        self.imageIds = imageIds

    def add_image_id(self, imageId):
        self.imageIds.append(imageId)

    # def __init__(self, imageFace):
    #     self.id = uuid.uuid4()
    #
    #     # 最大的imageFace的概率
    #     self.prob = imageFace.prob
    #
    #     # 最大的imageFace的向量值
    #     self.data = imageFace.data
    #
    #     # 这个人包含的imageFace的图片列表
    #     self.imageIds = []
    #     self.imageIds.append(imageFace.imageId)

    def distance(self, imageFace):
        # logger.info("【distance】face.data_type = %s", type(imageFace.data))
        # logger.info("【distance】face.data = %s", imageFace.data)
        #
        # logger.info("【distance】self.data_type = %s", type(self.data))
        # logger.info("【distance】self.data = %s", self.data)

        # distance = (self.data[0] - imageFace.data[0]).norm().item()
        distance = np.linalg.norm(np.array(self.data[0]) - np.array(imageFace.data[0]))
        logging.info("两张人脸的欧式距离为：%.2f", distance)
        return distance

    # 判断指定的imageFace是否属于当前这个人
    # def isSameOne(self,imageFace):
    #     if match_faces(self.data, imageFace.data, 1):
    #         self.imageIds.append(imageFace.imageId)
    #         if self.prob < imageFace.prob:
    #             self.prob = imageFace.prob
    #             self.data = imageFace.data
    #         return True
    #     else:
    #         return False

    def __lt__(self, other):
        return self.prob < other.prob

    # def __repr__(self):
    #     return f"id is {self.id} and prob is {self.prob} and data is {self.data} and imageIds is {self.imageIds}"

    def to_dict(self):
        return {
            "id": self.id,
            "prob": self.prob,
            "data": self.data,
            "imageIds": self.imageIds
        }


# # 计算人脸特征向量间的欧氏距离，设置阈值，判断是否为同一张人脸
# def match_faces(faces_emb, known_faces_emb, threshold):
#     result = False
#     distance = (known_faces_emb[0] - faces_emb[0]).norm().item()
#     logging.info("两张人脸的欧式距离为：%.2f", distance)
#
#     if distance < threshold:
#         result = True
#     return result


# 判断一张照片中是否有人脸，并且按照概率值，返回前3个人脸信息
def has_face(probs, face_tensors, face_prob_threshold, image_id):
    result = []
    result_test = []
    if probs is None or len(probs) == 0:
        return result

    for index in range(0, len(probs)):
        prob = probs[index]
        if prob is None or prob < face_prob_threshold:
            continue

        face_tensor = [face_tensors[index]]
        # 将多个单独的人脸张量堆叠成一个批量张量。
        aligned = torch.stack(face_tensor).to(device)

        # 将对齐后的人脸张量输入预训练模型resnet，提取每张人脸的特征向量
        face_emb = resnet(aligned).detach().cpu()

        # logger.debug("【HasFace】face_emb shape = %s", face_emb.shape)
        # logger.debug("【HasFace】face_emb = %s", face_emb)

        # embeddings = face_emb.detach().cpu().numpy()
        # logger.debug("【HasFace】embeddings = %s", embeddings)

        image_face = ImageFace(probs[index], face_emb.tolist(), index, image_id)
        result.append(image_face)
        result_test.append(image_face.to_dict())

    if len(result) > 0:
        sorted_list = sorted(result, reverse=True)
        result = sorted_list[:10]

    # logger.info("【HasFace】final_face_result = %s, result_length = %s", result_test, len(result_test))

    return result


persons = []


def update_db_person(index_name, person):
    return True


def write_person_file(folder_name, json_content, file_name):
    person_url = os.path.join(os.getcwd(), "data/person/" + folder_name)
    # logger.info("【WritePerson】write person file start. person_url = %s", person_url)

    # 判断folder是否存在，不存在就创建
    if not os.path.exists(person_url):
        # logger.info("【WritePerson】person_url not exist. folder_name = %s", folder_name)
        os.makedirs(person_url)

    final_file_path = person_url + "/" + file_name
    # logger.info("【WritePerson】write person file start. final_file_path = %s", final_file_path)

    with open(final_file_path, "w") as file:
        file.write(json_content)

    logger.info("【WritePerson】write person file finish. final_file_path = %s", final_file_path)


def get_person_from_db(index_name):
    result = []

    face = ImageFace(1.0, ["0.12", "0.23"], 1, "person.png")
    result.append(Person(face))

    return result


def get_person_from_disk(folder_name):
    result = []

    person_url = os.path.join(os.getcwd(), "data/person/" + folder_name)
    if not os.path.exists(person_url) or not os.path.isdir(person_url):
        return result

    files = os.listdir(person_url)

    # 过滤出所有的 JSON 文件
    json_files = [file for file in files if file.endswith('.json')]

    if len(json_files) == 0:
        return result

    for json_file in json_files:
        file_path = os.path.join(person_url, json_file)
        logger.info("【GetPersonFromDisk】file_path = %s", file_path)
        with open(file_path, 'r', encoding='utf-8') as file:
            try:
                # 读取并解析 JSON 文件
                json_file_content = json.load(file)

                # 如果是列表，遍历创建 Response 对象
                if isinstance(json_file_content, list):
                    for item in json_file_content:
                        result.append(Person(item.get("id"), item.get("prob"), item.get("data"), item.get("imageId"), item.get("imageIds")))
                # 如果是单个对象，直接创建 Response 对象
                elif isinstance(json_file_content, dict):
                    result.append(Person(json_file_content.get("id"), json_file_content.get("prob"), json_file_content.get("data"), json_file_content.get("imageId"), json_file_content.get("imageIds")))
            except json.JSONDecodeError as e:
                logger.info("文件 %s 解析失败: %s", json_file, e)

    return result


def process_face(image, application_name, index_name, face_threshold, image_id):
    start = time.time() * 1000
    face_tensors, face_probs = mtcnn(
        image
        , return_prob=True
        # , select_largest=False
        # ,save_path="/root/imageSearch/data/image/face/"+png_file
    )
    logger.info("【process_face】mtcnn process finish. cost time = %s", time.time() * 1000 - start)

    box_start_time = time.time() * 1000
    box_tensors, box_probs = mtcnn.detect(image)
    logger.info("【process_face】mtcnn detect finish. cost_time = %s", time.time() * 1000 - box_start_time)

    face_result = []

    # 判断当前照片中是否存在人脸，并返回相似度最高的10个人脸
    has_face_start_time = time.time() * 1000
    has_face_result = has_face(face_probs, face_tensors, face_threshold, image_id)
    logger.info("【process_face】has_face_result_length = %s, image_id = %s", len(has_face_result), image_id)
    logger.info("【process_face】has_face_result finish, cost_time = %s", time.time() * 1000 - has_face_start_time)

    if has_face_result is None or len(has_face_result) == 0:
        logging.info("【process_face】%s 不包含人脸，耗时 %f", image_id, time.time() - start)
        return face_result

    # logger.debug("【UploadImage】has_face_result = %s", has_face_result)

    business_start_time = time.time() * 1000

    logger.info("【process_face】%s 包含人脸,耗时 %f ,包含 %s 个人脸.", image_id, time.time() - start, len(has_face_result))
    for face in has_face_result:
        '''
            - 1、判断当前 index_name 下是否包含人的信息
                - 1.1、为空，说明当前index_name下没有人，则创建新的人
                - 1.1、不为空，说明当前index_name下有人，则需要继续处理：
                    - 1.1.1、for(Person person):
                        isSameOne(target, person):
                            - true:
                                - < 0.7?
                                    - true
                                        - 不用继续遍历person，直接更新当前人
                                    - false
                                        - 继续遍历person，记录最小的instance，将这个person更新，其他都不更新
                            = false:
                                - continue
        '''
        try:
            # 从数据库查询当前索引包含的person列表
            folder_name = base_64_encode(get_vector_index_name(application_name, index_name))
            logger.info("【process_face】get person from disk start. folder_name = %s, image_id = %s", folder_name,  image_id)

            person_list = get_person_from_disk(folder_name)
            logger.info("【process_face】get person from disk finish. person_list_size = %s, image_id = %s", len(person_list), image_id)

            if len(person_list) == 0:
                create_new_person(box_tensors, face, face_result, folder_name, image_id)
            else:
                compare_person_list = []
                has_find_owner = False
                # 3、当前index下有人，则需要继续处理
                for person in person_list:
                    # 4、判断当前这张脸和当前这个人是否匹配
                    distance = person.distance(face)
                    logger.info("【process_face】distance = %s, person_id = %s, person_prob = %s, face_prob = %s, image_id = %s", distance, person.id, person.prob, face.prob, image_id)
                    if distance > 0.8:
                        # 说明不是同一个人
                        logger.info("【process_face】distance is too large, distance = %s, person_id = %s, person_prob = %s, face_prob = %s, image_id = %s", distance, person.id, person.prob, face.prob, image_id)
                        continue
                    if distance <= 0.3:
                        logger.info("【process_face】distance <= 0.3, person_id = %s, person_prob = %s, face_prob = %s, image_id = %s", person.id, person.prob, face.prob, image_id)
                        compare_person_list.clear()
                        has_find_owner = True
                        # 直接替换，并break
                        if face.prob > person.prob:
                            logger.info("【process_face】distance <= 0.3, face.prob > person.prob, person_id = %s, person_prob = %s, face_prob = %s, image_id = %s", person.id, person.prob, face.prob, image_id)
                            # 替换
                            setattr(person, "prob", face.prob)
                            setattr(person, "data", face.data)
                            setattr(person, "imageId", image_id)
                            person.add_image_id(image_id)

                            person_json = json.dumps(person.to_dict(), indent=4)
                            write_person_file(folder_name, person_json, person.id + ".json")

                            face_result.append(FaceResponse(person.id, box_tensors[face.index].tolist(), face.prob, face.data[0]).to_dict())
                        else:
                            logger.info("【process_face】distance <= 0.3, face.prob <= person.prob, person_id = %s, person_prob = %s, face_prob = %s, image_id = %s", person.id, person.prob, face.prob, image_id)

                            person.add_image_id(image_id)

                            person_json = json.dumps(person.to_dict(), indent=4)
                            write_person_file(folder_name, person_json, person.id + ".json")

                            face_result.append(FaceResponse(person.id, box_tensors[face.index].tolist(), face.prob, face.data[0]).to_dict())
                        break
                    else:
                        # 不确定是同一个人
                        compare_person_list.append(person)

                if has_find_owner:
                    logger.info("【process_face】has_find_owner is true. image_id = %s", image_id)
                    continue

                # 在不确定的人中，选择相似度最大的人，然后判断这张脸和这个人的关系
                if len(compare_person_list) == 0:
                    logger.info("【process_face】compare_person_list is empty, add new person. face_prob = %s, image_id = %s", face.prob, image_id)
                    # 是一个新的人
                    create_new_person(box_tensors, face, face_result, folder_name, image_id)
                else:
                    sorted_person_list = sorted(compare_person_list, reverse=True)
                    max_person = sorted_person_list[0]
                    logger.info("【process_face】sorted_person_list is not empty, max_person_id = %s, max_person_prob = %s, face_prob = %s, image_id = %s", max_person.id, max_person.prob, face.prob, image_id)
                    if face.prob > max_person.prob:
                        logger.info("【process_face】face_prob > max_person_prob, person_id = %s, person_prob = %s, face_prob = %s, image_id = %s", max_person.id, max_person.prob, face.prob, image_id)
                        # 替换人脸
                        setattr(max_person, "prob", face.prob)
                        setattr(max_person, "data", face.data)
                        setattr(max_person, "image_id", image_id)
                        max_person.add_image_id(image_id)

                        person_json = json.dumps(max_person.to_dict(), indent=4)
                        write_person_file(folder_name, person_json, max_person.id + ".json")

                        face_result.append(FaceResponse(max_person.id, box_tensors[face.index].tolist(), face.prob, face.data[0]).to_dict())
                        continue
                    else:
                        max_person.add_image_id(image_id)

                        person_json = json.dumps(max_person.to_dict(), indent=4)
                        write_person_file(folder_name, person_json, max_person.id + ".json")

                        face_result.append(FaceResponse(max_person.id, box_tensors[face.index].tolist(), face.prob, face.data[0]).to_dict())
                        logger.info("【process_face】face_prob <= max_person_prob, person_id = %s, person_prob = %s, face_prob = %s, image_id = %s", max_person.id, max_person.prob, face.prob, image_id)
        except Exception as e:
            logger.error("【process_face】process face error. image_id = %s, error = %s", image_id, e)

    logger.info("【process_face】process face end. cost time = %s", time.time() * 1000 - business_start_time)
    return face_result


def create_new_person(box_tensors, face, face_result, folder_name, image_id):
    # 2、当前index下没有人，则创建新的人
    person_id = str(uuid.uuid4())

    logger.info("【process_face】create new person start. person_id = %s, image_id = %s", person_id, image_id)
    new_person = Person(person_id, face.prob, face.data, image_id, [image_id])
    # 写入文件，文件名为person_id.json
    person_json = json.dumps(new_person.to_dict(), indent=4)
    write_person_file(folder_name, person_json, person_id + ".json")
    # face_result.append(FaceResponse(new_person.id, new_person.prob, None, new_person.data))
    face_result.append(FaceResponse(new_person.id, box_tensors[face.index].tolist(), new_person.prob, new_person.data[0]).to_dict())

